package com.stress.testing.plugin.apache.kafka.interceptor;

import com.stress.testing.common.StressTesting;
import com.stress.testing.common.TraceContext;
import com.stress.testing.common.TraceContextHolder;
import com.stress.testing.common.interceptor.AbstractInterceptor;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.header.internals.RecordHeaders;

import java.lang.reflect.Field;
import java.lang.reflect.Method;

import static com.stress.testing.common.conf.PluginConfigKey.DEFAULT_PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX;
import static com.stress.testing.common.conf.PluginConfigKey.PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX_KEY;

public class KafkaProducerSendInterceptor extends AbstractInterceptor {
    private Field topicField;
    private Field headersField;

    public KafkaProducerSendInterceptor() {
        Class<ProducerRecord> recordClass = ProducerRecord.class;
        try {
            this.topicField = recordClass.getDeclaredField("topic");
        } catch (NoSuchFieldException e) {
            throw new IllegalArgumentException(
                    "org.apache.kafka.clients.producer.ProducerRecord has no field named topic", e);
        }
        try {
            this.headersField = recordClass.getDeclaredField("headers");
        } catch (NoSuchFieldException e) {
            // 版本较低的是没有这个字段的
        }
    }

    @Override
    public void before(Method method, Object[] args, Object zuper, Object thiz) {
        if (!StressTesting.enable(getAgentConfig(), getPluginConfig(), TraceContextHolder.getTraceContext())) {
            return;
        }

        if (args == null || args.length != 2) {
            return;
        }

        Object arg = args[0];
        if (arg == null || !(arg instanceof ProducerRecord)) {
            return;
        }

        ProducerRecord record = (ProducerRecord) arg;
        changeTopic(record);
        fillHeaders(record);
    }

    private void fillHeaders(ProducerRecord record) {
        if (headersField == null) {
            return;
        }

        Headers headers = record.headers();
        if (headers == null) {
            headers = new RecordHeaders();

            headersField.setAccessible(true);
            try {
                headersField.set(record, headers);
            } catch (IllegalAccessException e) {
                // warning log
            } finally {
                headersField.setAccessible(true);
            }
        }

        TraceContext traceContext = TraceContextHolder.getTraceContext();
        if (traceContext == null) {
            return;
        }

        // 添加链路信息
        for (String name : traceContext.propertyNames()) {
            headers.add(name, traceContext.getProperty(name).getBytes());
        }
    }

    private void changeTopic(ProducerRecord record) {
        String topicOrigin = record.topic();
        String topicPrefix = getPluginConfig().get(PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX_KEY,
                DEFAULT_PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX);

        if (topicOrigin == null || topicOrigin.startsWith(topicPrefix)) {
            return;
        }

        String newTopic = topicPrefix + topicOrigin;

        topicField.setAccessible(true);
        try {
            topicField.set(record, newTopic);
        } catch (IllegalAccessException e) {
            // warning log
        } finally {
            topicField.setAccessible(false);
        }
    }
}
